package com.yunchao.library.socket

import android.annotation.SuppressLint
import android.content.Context
import android.net.*
import android.os.Build
import com.yunchao.library.socket.config.ConnectionState
import com.yunchao.library.socket.config.Message
import com.yunchao.library.socket.config.MessageContent
import com.yunchao.library.socket.config.MessageHeader
import com.yunchao.library.socket.config.TcpConfig
import com.yunchao.library.socket.enum.ConnectionEvent
import com.yunchao.library.socket.enum.TransportType
import com.yunchao.library.socket.exception.SocketClientException
import com.yunchao.library.socket.metrics.BasicSocketMetrics
import com.yunchao.library.socket.metrics.SocketMetrics
import kotlinx.coroutines.*
import kotlinx.coroutines.channels.BufferOverflow
import kotlinx.coroutines.flow.*
import kotlinx.coroutines.sync.*
import java.io.*
import java.net.InetSocketAddress
import java.net.Socket
import java.util.concurrent.ConcurrentLinkedQueue
import java.util.concurrent.atomic.AtomicBoolean
import java.util.concurrent.atomic.AtomicInteger
import kotlin.math.*

/**
 * 高性能TCP Socket客户端实现
 */
class TcpSocketClient(
    private val config: TcpConfig,
    private val externalScope: CoroutineScope,
    private val context: Context,
    private val socketFactory: () -> Socket = { Socket() },
    private val metrics: SocketMetrics = BasicSocketMetrics().apply { reset() }
) : Closeable, CoroutineScope by externalScope {

    // 单线程发送调度器
    @OptIn(ExperimentalCoroutinesApi::class)
    private val sendDispatcher = Dispatchers.IO.limitedParallelism(1)

    // region 状态管理
    private val _connectionState = MutableStateFlow<ConnectionState>(ConnectionState.Disconnected)
    val connectionState: StateFlow<ConnectionState> = _connectionState.asStateFlow()
    // endregion

    // region 核心组件
    private val closed = AtomicBoolean(false)
    private val job = Job()
    private val ioScope = CoroutineScope(Dispatchers.IO + job)
    private val connectionLock = Mutex()
    private val bufferPool = SmartBufferPool(listOf(1024, 2048, 4096, 8192, 16384))

    private var heartbeatJob: Job? = null
    private var reconnectJob: Job? = null
    private var reconnectAttempts = AtomicInteger(0)
    private var networkCallback: ConnectivityManager.NetworkCallback? = null

    private val _incomingMessages = MutableSharedFlow<Message>(
        extraBufferCapacity = config.messageBufferCapacity,
        onBufferOverflow = BufferOverflow.SUSPEND
    )
    val incomingMessages: SharedFlow<Message> = _incomingMessages.asSharedFlow()
    // endregion

    init {
        registerNetworkCallback()
    }

    // region 公开API
    fun start() {
        checkNotDisposed()
        ioScope.launch { connect() }
    }

    suspend fun sendMessage(
        direction: Int,
        transport: TransportType,
        bodyBytes: ByteArray,
        isPacket: Int = 0,
        indicator: Int = 0,
        sequence: Long = 0
    ) {
        checkConnected()
        withContext(sendDispatcher) {
            try {
                internalSend(direction, transport, bodyBytes, isPacket, indicator, sequence)
            } catch (e: Exception) {
                handleConnectionError(e)
            }
        }
    }

    override fun close() {
        if (closed.compareAndSet(false, true)) {
            ioScope.launch {
                disconnectInternal()
                cleanupResources()
                job.cancel()
            }
            metrics.recordConnectionEvent(ConnectionEvent.Disconnected)
        }
    }
    // endregion

    // region 连接管理
    private suspend fun connect() = connectionLock.withLock {
        if (shouldSkipConnect()) return@withLock

        updateState(ConnectionState.Connecting)
        metrics.recordConnectionEvent(ConnectionEvent.Connecting)

        try {
            disconnectInternal()
            val socket = socketFactory().apply {
                connect(InetSocketAddress(config.host, config.port), config.connectTimeout)
                keepAlive = true
                tcpNoDelay = true
                soTimeout = 0
            }
            handleSuccessfulConnection(socket)
        } catch (e: Exception) {
            handleConnectionError(e)
        }
    }

    private suspend fun disconnect() = connectionLock.withLock {
        if (shouldSkipDisconnect()) return@withLock
        disconnectInternal()
    }

    private fun handleSuccessfulConnection(socket: Socket) {
        if (!closed.compareAndSet(true, false)) {
            closed.set(false) // 确保状态正确
        }
//        startHeartbeat()
        updateState(ConnectionState.Connected(socket))
        reconnectAttempts.set(0)
        metrics.recordConnectionEvent(ConnectionEvent.Connected(config.host, config.port))
        receiveMessages(socket)
    }
    // endregion

    // region 网络监听
    @SuppressLint("ObsoleteSdkInt")
    private fun registerNetworkCallback() {
        if (Build.VERSION.SDK_INT < Build.VERSION_CODES.N) return

        val cm = context.getSystemService(Context.CONNECTIVITY_SERVICE) as ConnectivityManager
        val request = NetworkRequest.Builder()
            .addCapability(NetworkCapabilities.NET_CAPABILITY_INTERNET)
            .build()

        networkCallback = object : ConnectivityManager.NetworkCallback() {
            override fun onAvailable(network: Network) {
                ioScope.launch {
                    disconnect()
                    delay(1000)
                    connect()
                }
            }

            override fun onLost(network: Network) {
                ioScope.launch { disconnect() }
            }
        }.also {
            cm.registerNetworkCallback(request, it)
        }
    }
    // endregion

    // region 消息处理
    private fun internalSend(
        direction: Int,
        transport: TransportType,
        bodyBytes: ByteArray,
        isPacket: Int,
        indicator: Int,
        sequence: Long
    ) {
        val isChunked = bodyBytes.size > config.largeMessageThreshold
        val header = MessageHeader(
            bodyLength = bodyBytes.size,
            direction = direction,
            isPacket = isPacket,
            indicator = indicator,
            sequence = sequence,
            transport = transport,
            isChunked = isChunked
        )
        metrics.recordMessageSent(header, bodyBytes)

        val socket = (connectionState.value as ConnectionState.Connected).socket
        val outputStream = socket.getOutputStream()

        outputStream.write(buildHeaderBytes(header))

        when {
            isChunked -> sendInChunks(outputStream, bodyBytes)
            else -> outputStream.write(bodyBytes)
        }
        outputStream.flush()
    }

    private fun receiveMessages(socket: Socket) = ioScope.launch {
        val inputStream = socket.getInputStream()
        val headerBuffer = ByteArray(14)

        while (isActive && isConnected()) {
            try {
                readFully(inputStream, headerBuffer)
                val header = parseHeader(headerBuffer)

                val message = if (header.isChunked) {
                    receiveChunkedMessage(inputStream, header)
                } else {
                    Message(
                        header,
                        MessageContent.SingleMessage(readBody(inputStream, header.bodyLength))
                    )
                }
                _incomingMessages.emit(message).also {
                    metrics.recordMessageReceived(message)
                }
            } catch (e: Exception) {
                if (isConnected()) handleConnectionError(e)
            }
        }
    }
    // endregion

    // region 协议处理

    // ================== 协议解析 ==================
    /**
     * 解析消息头
     */
    private fun parseHeader(headerBytes: ByteArray): MessageHeader {
        if (headerBytes.size != 14) {
            throw SocketClientException.ProtocolError("Invalid header length")
        }
        DataInputStream(ByteArrayInputStream(headerBytes)).use { dis ->
            val firstFourBytes = dis.readInt()
            val direction = (firstFourBytes shr 31) and 1
            val isPacket = (firstFourBytes shr 30) and 1
            val indicator = (firstFourBytes shr 28) and 0x3
            val bodyLength = firstFourBytes and 0x0FFFFFFF
            val sequence = dis.readLong()
            val byte1 = dis.readByte().toInt() and 0xFF
            val byte2 = dis.readByte().toInt() and 0xFF
            val transportCode = (byte1 shl 8) or byte2
            val transport = TransportType.fromCode(transportCode)
            val isChunked = bodyLength > config.largeMessageThreshold
            return MessageHeader(
                bodyLength = bodyLength,
                direction = direction,
                isPacket = isPacket,
                indicator = indicator,
                sequence = sequence,
                transport = transport,
                isChunked = isChunked
            )
        }
    }

    /**
     * 构建消息头字节数组
     */
    private fun buildHeaderBytes(header: MessageHeader): ByteArray {
        var firstFourBytes = header.bodyLength
        firstFourBytes = firstFourBytes or (header.direction shl 31)
        firstFourBytes = firstFourBytes or (header.isPacket shl 30)
        firstFourBytes = firstFourBytes or (header.indicator shl 28)

        ByteArrayOutputStream(14).use { baos ->
            DataOutputStream(baos).use { dos ->
                dos.writeInt(firstFourBytes)
                dos.writeLong(header.sequence)
                val transportCode = header.transport.code
                dos.writeByte(transportCode ushr 8)
                dos.writeByte(transportCode and 0xFF)
                return baos.toByteArray()
            }
        }
    }
    // endregion

    // region 分块传输
    private fun sendInChunks(outputStream: OutputStream, data: ByteArray) {
        val buffer = bufferPool.acquire(config.chunkSize)
        try {
            var offset = 0
            while (offset < data.size && isConnected()) {
                val chunkSize = min(buffer.size, data.size - offset)
                System.arraycopy(data, offset, buffer, 0, chunkSize)
                outputStream.write(buffer, 0, chunkSize)

                metrics.recordChunkSent(offset, chunkSize, data)
                offset += chunkSize
            }
        } finally {
            bufferPool.release(buffer)
        }
    }

    private fun receiveChunkedMessage(inputStream: InputStream, header: MessageHeader): Message {
        val chunks = mutableListOf<ByteArray>()
        var remaining = header.bodyLength
        val buffer = bufferPool.acquire(config.chunkSize)

        try {
            while (remaining > 0 && isActive) {
                val readSize = min(buffer.size, remaining)
                readFully(inputStream, buffer, readSize)
                chunks.add(buffer.copyOf(readSize))
                remaining -= readSize
            }

            return Message(
                header = header,
                content = MessageContent.ChunkedMessage(
                    chunks = chunks,
                    totalSize = header.bodyLength
                )
            )
        } finally {
            bufferPool.release(buffer)
        }
    }
    // endregion

    // region 工具方法
    private fun updateState(newState: ConnectionState) {
        _connectionState.value = newState
    }

    private fun checkConnected() {
        if (closed.get() || !isConnected()) {
            throw SocketClientException.ConnectionFailed()
        }
    }

    private fun checkNotDisposed() {
        if (closed.get()) throw SocketClientException.Disposed()
    }

// 判断连接状态是否为已连接
    private fun isConnected() = connectionState.value is ConnectionState.Connected

    // 判断是否跳过连接
    private fun shouldSkipConnect(): Boolean {
        // 如果连接已关闭，则返回true
        return closed.get() || connectionState.value.let {
            // 如果连接状态为已连接或正在连接，则返回true
            it is ConnectionState.Connected || it is ConnectionState.Connecting
        }
    }

    private fun shouldSkipDisconnect(): Boolean {
        return closed.get() || connectionState.value !is ConnectionState.Connected
    }

    private fun readFully(input: InputStream, buffer: ByteArray, length: Int = buffer.size) {
        var offset = 0
        while (offset < length) {
            val read = input.read(buffer, offset, length - offset)
            if (read == -1) throw EOFException("Unexpected end of stream")
            offset += read
        }
    }

    private fun readBody(input: InputStream, length: Int): ByteArray {
        return ByteArray(length).also {
            readFully(input, it)
        }
    }
    // endregion

    // region 心跳机制
    fun startHeartbeat() {
        heartbeatJob?.cancel()
        heartbeatJob = ioScope.launch {
            while (isActive && isConnected()) {
                delay(config.heartbeatInterval)
                try {
                    internalSend(
                        direction = 1,
                        transport = TransportType.HEARTBEAT,
                        bodyBytes = byteArrayOf(),
                        isPacket = 0,
                        indicator = 0,
                        sequence = 0
                    )
                } catch (e: Exception) {
                    handleConnectionError(e)
                }
            }
        }
    }
    // endregion

    // region 错误处理
    // 处理连接错误
    private fun handleConnectionError(error: Throwable) {
        // 如果已经关闭，则直接返回
        if (closed.get()) return

        // 在ioScope中执行
        ioScope.launch {
            // 获取连接锁
            connectionLock.withLock {
                // 断开连接
                disconnectInternal()
                // 根据不同的情况处理错误
                when {
                    // 如果不自动重连，则更新状态为错误，并记录连接事件
                    !config.autoReconnect -> {
                        updateState(ConnectionState.Error(error))
                        metrics.recordConnectionEvent(ConnectionEvent.Error(error))
                    }

                    // 如果重连次数超过最大重连次数，则更新状态为错误，并记录连接事件
                    reconnectAttempts.get() >= config.maxReconnectAttempts -> {
                        updateState(ConnectionState.Error(SocketClientException.MaxReconnectAttemptsReached()))
                        metrics.recordConnectionEvent(ConnectionEvent.Error(SocketClientException.MaxReconnectAttemptsReached()))
                    }

                    // 否则，调度重连
                    else -> scheduleReconnect()
                }
            }
        }
    }

    // 定义一个函数，用于重新连接
    private fun scheduleReconnect() {
        // 增加重新连接尝试次数
        reconnectAttempts.incrementAndGet()
        // 计算延迟时间，取最小值
        val delayMs = min(
            // 计算延迟时间，使用指数增长
            config.reconnectBaseDelay * 2.0.pow(reconnectAttempts.get().toDouble() - 1).toLong(),
            // 最大延迟时间
            config.maxReconnectDelay
        )
        // 记录连接事件
        metrics.recordConnectionEvent(
            ConnectionEvent.Reconnecting(
                // 重新连接尝试次数
                reconnectAttempts.get(),
                // 延迟时间
                delayMs
            )
        )
        // 启动一个延迟任务，延迟时间为delayMs，执行connect()函数
        reconnectJob = ioScope.launch {
            delay(delayMs)
            connect()
        }
    }
    // endregion

    // region 资源管理
    @SuppressLint("ObsoleteSdkInt")
    private fun cleanupResources() {
        bufferPool.clear()
        if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.N) {
            (context.getSystemService(Context.CONNECTIVITY_SERVICE) as ConnectivityManager)
                .unregisterNetworkCallback(networkCallback!!)
        }
    }

    private fun disconnectInternal() {
        (connectionState.value as? ConnectionState.Connected)?.socket?.close()
        updateState(ConnectionState.Disconnected)
        heartbeatJob?.cancel()
        reconnectJob?.cancel()
    }
    // endregion

    // region 内部类
    private class SmartBufferPool(private val sizes: List<Int>) {
        private val pools = sizes.associateWith { ConcurrentLinkedQueue<ByteArray>() }

        fun acquire(size: Int): ByteArray {
            val suitableSize = sizes.firstOrNull { it >= size } ?: size
            return pools[suitableSize]?.poll() ?: ByteArray(suitableSize)
        }

        fun release(buffer: ByteArray) {
            pools[buffer.size]?.offer(buffer)
        }

        fun clear() {
            pools.values.forEach { it.clear() }
        }
    }
    // endregion
}
